/*========================== begin_copyright_notice ============================

Copyright (C) 2021 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#include "common/LLVMWarningsPush.hpp"
#include <llvm/ADT/SetVector.h>
#include <llvm/IR/InstIterator.h>
#include <llvm/Transforms/Utils/Cloning.h>
#include <llvm/Transforms/IPO.h>
#include "common/LLVMWarningsPop.hpp"

#include <common/LLVMUtils.h>
#include <common/ModuleSplitter.h>
#include <Compiler/CodeGenPublic.h>
#include "Compiler/CISACodeGen/OpenCLKernelCodeGen.hpp"

namespace IGC {
KernelModuleSplitter::KernelModuleSplitter(IGC::OpenCLProgramContext &oclContext, llvm::Module &module)
    : _oclContext(oclContext), _originalModule(module), _splittedModule(nullptr) {}

KernelModuleSplitter::~KernelModuleSplitter() { restoreOclContextModule(); }

void KernelModuleSplitter::splitModuleForKernel(const llvm::Function *kernelF) {
  using namespace llvm;
  IGC_ASSERT_EXIT_MESSAGE(kernelF != nullptr, "Cannot split for null function!");

  std::vector<const Function *> workqueue;
  SetVector<const GlobalValue *> GVs;

  // add all functions called by the kernel, recursively
  // start with the kernel...
  GVs.insert(kernelF);
  workqueue.push_back(kernelF);

  // and for all called functions...
  while (!workqueue.empty()) {
    const Function *F = workqueue.back();
    workqueue.pop_back();

    for (const auto &I : instructions(F)) {
      if (const CallBase *CB = dyn_cast<CallBase>(&I)) {
        if (const Function *CF = CB->getCalledFunction()) {
          if (CF->isDeclaration() || GVs.count(CF))
            continue;

          // add only defined ones and rerun for their's calls ...
          GVs.insert(CF);
          workqueue.push_back(CF);
        }
      }
    }
  }

  // add all globals - it's easier to let them be removed later than search for them here
  for (auto &GV : _originalModule.globals()) {
    GVs.insert(&GV);
  }

  // create new module with selected globals and functions
  ValueToValueMapTy VMap;
  std::unique_ptr<Module> kernelM =
      CloneModule(_originalModule, VMap, [&](const GlobalValue *GV) { return GVs.count(GV); });
  IGC_ASSERT_EXIT_MESSAGE(kernelM, "Cloning module failed!");

  // Do cleanup.
  IGC::IGCPassManager mpm(&_oclContext, "CleanupAfterModuleSplitting");
  mpm.add(createGlobalDCEPass());           // Delete unreachable globals.
  mpm.add(createStripDeadDebugInfoPass());  // Remove dead debug info.
  mpm.add(createStripDeadPrototypesPass()); // Remove dead func decls.

  mpm.run(*kernelM.get());
  _splittedModule = std::move(kernelM);
}

void KernelModuleSplitter::retry() {
  if (_splittedModule) {
    restoreOclContextModule();
    delete _splittedModule.release();
  }
}

void KernelModuleSplitter::restoreOclContextModule() {
  if (_splittedModule) {
    _oclContext.clearMD();
    _oclContext.setModule(&_originalModule);
  }
}

void KernelModuleSplitter::setSplittedModuleInOCLContext() {
  if (_splittedModule) {
    _oclContext.clearMD();
    _oclContext.setModule(_splittedModule.get());
  }
}
} // namespace IGC